#!/usr/bin/env python3
"""
简单的已知明文攻击方案
通过发送已知内容并捕获加密结果来建立映射表
"""

import json
import base64
import time
from datetime import datetime

class PlaintextAttackHelper:
    def __init__(self):
        # 已知的映射
        self.mapping = {
            0xa1: ord('m'),
            0xf3: ord('o'),
            0x6b: ord('d'),
            0x31: ord('e'),
            0xb0: ord('l'),
            0x69: ord('"'),
            0x01: ord(':'),
            0x90: ord(' ')
        }
        
        # 需要找到映射的字符
        self.missing_chars = []
        for i in range(32, 127):  # ASCII可打印字符
            if i not in self.mapping.values():
                self.missing_chars.append(chr(i))
        
        print(f"已知映射: {len(self.mapping)} 个")
        print(f"待映射字符: {len(self.missing_chars)} 个")
    
    def generate_test_requests(self):
        """生成测试请求"""
        requests = []
        
        # 1. 数字测试
        requests.append({
            "questionText": "Test digits: 0123456789",
            "timestamp": str(int(time.time())),
            "requestId": "test_digits"
        })
        
        # 2. 大写字母测试  
        requests.append({
            "questionText": "UPPERCASE: ABCDEFGHIJKLMNOPQRSTUVWXYZ",
            "userId": "TEST_USER",
            "requestId": "test_upper"
        })
        
        # 3. 小写字母测试（跳过已知的）
        lowercase = "abcdefghijklmnopqrstuvwxyz"
        unknown_lower = ''.join(c for c in lowercase if ord(c) not in self.mapping.values())
        requests.append({
            "questionText": f"lowercase test: {unknown_lower}",
            "sessionId": "test_session_123",
            "requestId": "test_lower"
        })
        
        # 4. 特殊字符测试
        special = "!@#$%^&*()_+-=[]{}|;':<>?,./"
        unknown_special = ''.join(c for c in special if ord(c) not in self.mapping.values())
        requests.append({
            "questionText": f"Special chars: {unknown_special}",
            "type": "special_test",
            "requestId": "test_special"
        })
        
        # 5. 混合测试
        requests.append({
            "questionText": "Mixed: Hello123!@# World456$%^",
            "conversationId": "conv_test_789",
            "requestId": "test_mixed"
        })
        
        return requests
    
    def print_test_request(self, request):
        """打印测试请求供用户发送"""
        print("\n" + "="*60)
        print("请在IDEA中发送以下内容作为questionText:")
        print("-"*60)
        print(f"内容: {request['questionText']}")
        print(f"请求ID: {request['requestId']}")
        print("-"*60)
        print("完整JSON:")
        print(json.dumps(request, indent=2, ensure_ascii=False))
        print("="*60)
        
    def analyze_captured_data(self, plaintext_json, captured_encrypted):
        """
        分析捕获的加密数据
        plaintext_json: 发送的原始JSON
        captured_encrypted: 捕获的加密请求体
        """
        print("\n分析捕获的数据...")
        
        # 1. 解码自定义Base64
        charset = "!#%&()*,.@ABCDEFGHIJKLMNOPQRSTUVWXYZ^_abcdefghijklmnopqrstuvwxyz"
        std_base64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
        
        trans_table = str.maketrans(charset, std_base64)
        translated = captured_encrypted.translate(trans_table)
        
        # 添加padding
        padding = (4 - len(translated) % 4) % 4
        translated += '=' * padding
        
        # Base64解码
        try:
            decoded = base64.b64decode(translated)
        except Exception as e:
            print(f"Base64解码失败: {e}")
            return
        
        print(f"解码后长度: {len(decoded)} 字节")
        
        # 2. 查找已知文本在解码数据中的位置
        plaintext_bytes = plaintext_json.encode('utf-8')
        
        # 搜索已知的模式
        # 例如 "questionText": 在加密后应该是 69 ae 1d ... 69 01 90
        
        # 查找所有可能匹配的位置
        self.find_text_mapping(plaintext_bytes, decoded)
    
    def find_text_mapping(self, plaintext, ciphertext):
        """查找明文和密文的对应关系"""
        print("\n查找文本映射...")
        
        # 查找已知的模式来定位
        # "questionText" 应该映射到特定的密文
        
        # 搜索 "questionText" 的可能位置
        target = b"questionText"
        
        # 在密文中查找可能的位置
        # 已知 " 映射到 0x69, : 映射到 0x01
        
        new_mappings = {}
        
        # 尝试找到questionText的位置
        for i in range(len(ciphertext) - len(target)):
            # 检查是否可能是questionText的加密
            possible = True
            temp_mapping = {}
            
            for j, plain_byte in enumerate(target):
                cipher_byte = ciphertext[i + j]
                
                # 检查已知映射
                if plain_byte in self.mapping.values():
                    # 找到对应的密文字节
                    expected_cipher = None
                    for k, v in self.mapping.items():
                        if v == plain_byte:
                            expected_cipher = k
                            break
                    
                    if expected_cipher and cipher_byte != expected_cipher:
                        possible = False
                        break
                else:
                    # 记录新的映射
                    temp_mapping[cipher_byte] = plain_byte
            
            if possible and len(temp_mapping) > 0:
                print(f"\n可能在偏移 {i} 找到 'questionText'")
                print(f"新映射: {len(temp_mapping)} 个")
                
                # 添加新映射
                for k, v in temp_mapping.items():
                    if k not in new_mappings:
                        new_mappings[k] = v
                        print(f"  0x{k:02x} -> '{chr(v)}' (ASCII {v})")
        
        if new_mappings:
            print(f"\n✅ 找到 {len(new_mappings)} 个新映射!")
            
            # 更新映射表
            self.mapping.update(new_mappings)
            print(f"当前总映射: {len(self.mapping)} 个")
            
            # 保存当前映射表
            self.save_mapping()
        else:
            print("\n未找到新映射")
    
    def save_mapping(self):
        """保存当前映射表"""
        # 创建256字节的映射表
        table = bytearray(range(256))
        
        for k, v in self.mapping.items():
            table[k] = v
        
        filename = f'mapping_table_{len(self.mapping)}_chars.bin'
        with open(filename, 'wb') as f:
            f.write(table)
        
        print(f"映射表已保存到: {filename}")
        
        # 同时保存为JSON便于查看
        json_mapping = {f"0x{k:02x}": chr(v) if 32 <= v <= 126 else f"\\x{v:02x}" 
                       for k, v in self.mapping.items()}
        
        with open('mapping_progress.json', 'w') as f:
            json.dump({
                'mapping_count': len(self.mapping),
                'timestamp': datetime.now().isoformat(),
                'mappings': json_mapping
            }, f, indent=2)
    
    def test_current_mapping(self):
        """测试当前映射表"""
        if not os.path.exists('decoded_result.bin'):
            print("未找到decoded_result.bin")
            return
        
        print(f"\n使用当前映射表测试 ({len(self.mapping)} 个映射)...")
        
        with open('decoded_result.bin', 'rb') as f:
            encrypted = f.read()
        
        # 创建映射表
        table = bytearray(range(256))
        for k, v in self.mapping.items():
            table[k] = v
        
        # 解密
        decrypted = bytes(table[b] for b in encrypted)
        
        # 统计
        mapped_bytes = sum(1 for b in encrypted if b in self.mapping)
        print(f"已映射字节: {mapped_bytes}/{len(encrypted)} ({mapped_bytes/len(encrypted)*100:.1f}%)")
        
        # 保存结果
        with open(f'partial_decrypt_{len(self.mapping)}.txt', 'wb') as f:
            f.write(decrypted)
        
        # 显示预览
        preview = decrypted[:300].decode('utf-8', errors='ignore')
        print(f"\n预览:\n{preview}")

def main():
    """主流程"""
    import os
    
    print("=" * 60)
    print("已知明文攻击助手")
    print("=" * 60)
    
    attack = PlaintextAttackHelper()
    
    # 生成测试请求
    test_requests = attack.generate_test_requests()
    
    print("\n📝 使用说明:")
    print("1. 运行此脚本生成测试请求")
    print("2. 在IDEA中发送这些请求")
    print("3. 使用代理工具捕获加密的请求体")
    print("4. 将捕获的数据保存为 captured_1.txt, captured_2.txt 等")
    print("5. 再次运行此脚本进行分析")
    
    # 检查是否有捕获的数据
    captured_files = []
    for i in range(1, 10):
        if os.path.exists(f'captured_{i}.txt'):
            captured_files.append(f'captured_{i}.txt')
    
    if captured_files:
        print(f"\n找到 {len(captured_files)} 个捕获文件")
        
        for i, captured_file in enumerate(captured_files):
            print(f"\n分析 {captured_file}...")
            
            with open(captured_file, 'r') as f:
                captured_data = f.read().strip()
            
            # 对应的测试请求
            if i < len(test_requests):
                test_request = test_requests[i]
                plaintext = json.dumps(test_request, ensure_ascii=False)
                
                attack.analyze_captured_data(plaintext, captured_data)
        
        # 测试当前映射
        attack.test_current_mapping()
        
    else:
        print("\n未找到捕获文件，显示测试请求:")
        
        for i, request in enumerate(test_requests):
            attack.print_test_request(request)
            
            print(f"\n请将捕获的加密请求体保存为: captured_{i+1}.txt")
            input("按Enter继续下一个请求...")

if __name__ == "__main__":
    main()